/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.example.hrpc.rpc.client;

import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.channel.*;
import io.netty.channel.socket.SocketChannel;
import io.netty.handler.timeout.IdleStateHandler;
import org.example.hrpc.common.utils.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

import static org.example.hrpc.rpc.handler.Decoder.DECODER;
import static org.example.hrpc.rpc.handler.Encoder.ENCODER;

/**
 * Factory for creating {@link TransportClient}s by using createClient.
 * 创建TransportClient的工厂，工厂维护了
 */
public class TransportClientFactory implements Closeable {
  public final static TransportClientFactory INSTANCE = new TransportClientFactory();

  /** A simple data structure to track the pool of clients between two peer nodes.
   * TransportClientFactory内部还存在针对每个Socket地址的连接池ClientPool
   * ClientPool实际是由TransportClient的数组构成的
   * lock：通过对每个TransportClient分别采用不同的锁，降低并发情况下线程间对锁的争用，进而减少阻塞，提高并发度
   */
  private static class ClientPool {
    TransportClient[] clients;
    Object[] locks;

    ClientPool(int size) {
      clients = new TransportClient[size];
      locks = new Object[size];
      for (int i = 0; i < size; i++) {
        locks[i] = new Object();
      }
    }
  }

  private static final Logger logger = LoggerFactory.getLogger(TransportClientFactory.class);

  private final ConcurrentHashMap<SocketAddress, ClientPool> connectionPool;

  /** Random number generator for picking connections between peers. */
  private final Random rand;
  private final int numConnectionsPerPeer = 8;

  private final Class<? extends Channel> socketChannelClass;
  private EventLoopGroup workerGroup;
  private final PooledByteBufAllocator pooledAllocator;


  public TransportClientFactory() {
    /**
     * 针对每个Socket地址的连接池ClientPool的缓存：
     * 也就是说，每一个Socket地址可能有多个Transportclient对象，
     * 有点不对劲：
     */
    this.connectionPool = new ConcurrentHashMap<>();

//    this.numConnectionsPerPeer = conf.numConnectionsPerPeer();
    /**
     * 对Socket地址对应的连接池ClientPool中缓存的TransportClient进行随机选择，对每个连接做负载均衡
     */
    this.rand = new Random();

    /**
     * IO模式，默认值为NIO,还支持EPOLL
     */
    IOMode ioMode = IOMode.NIO;
    /**
     * 客户端Channel被创建时使用的类，通过ioMode来匹配，默认为NioSocketChannel, Spark还支持EpollEventLoopGroup
     */
    this.socketChannelClass = NettyUtils.getClientChannelClass(ioMode);
    int numThreads = Math.min(Runtime.getRuntime().availableProcessors() + 1, 32);
    this.workerGroup = NettyUtils.createEventLoop(ioMode, numThreads, "rpc-client");

    this.pooledAllocator = NettyUtils.getSharedPooledByteBufAllocator(false, false);
  }


  /**
   * Create a {@link TransportClient} connecting to the given remote host / port.
   *
   * We maintains an array of clients (size determined by spark.shuffle.io.numConnectionsPerPeer)
   * and randomly picks one to use. If no client was previously created in the randomly selected
   * spot, this function creates a new client and places it there.
   * 我们维护了clients数组，并随机选择一个使用。如果在随机位置上过去没有client被创建，这个函数创建一个新的client并放置在该数组位置上
   *
   * Concurrency: This method is safe to call from multiple threads.
   * 每个TransportClient实例只能和一个远端的RPC服务通信
   * 所以当一个组件想要和多个服务通信时，就需要持有多个TransportClient实例
   */
  public TransportClient createClient(String remoteHost, int remotePort)
      throws IOException, InterruptedException {
    // Get connection from the connection pool first.
    // If it is not found or not active, create a new one.
    // Use unresolved address here to avoid DNS resolution each time we creates a client.
    final InetSocketAddress unresolvedAddress =
      InetSocketAddress.createUnresolved(remoteHost, remotePort);

    // Create the ClientPool if we don't have it yet.
    ClientPool clientPool = connectionPool.get(unresolvedAddress);
    if (clientPool == null) {
      connectionPool.putIfAbsent(unresolvedAddress, new ClientPool(numConnectionsPerPeer));
      clientPool = connectionPool.get(unresolvedAddress);
    }
    /**
     * 随机选择一个TransportClient
     * 根据numConnectionPerPeer的大小（使用  spark.模块名.io.numConnectionsPerPeer 属性配置）
     */
    int clientIndex = rand.nextInt(numConnectionsPerPeer);
    TransportClient cachedClient = clientPool.clients[clientIndex];

    if (cachedClient != null && cachedClient.isActive()) {
      // Make sure that the channel will not timeout by updating the last use time of the
      // handler. Then check that the client is still alive, in case it timed out before
      // this code was able to update things.
      // 确保channel没有超时

      if (cachedClient.isActive()) {
        logger.trace("Returning cached connection to {}: {}",
          cachedClient.getSocketAddress(), cachedClient);
        return cachedClient;
      }
    }

    // If we reach here, we don't have an existing connection open. Let's create a new one.
    // Multiple threads might race here to create new connections. Keep only one of them active.
    // 当程序执行到这里时，表明没有可用Connection对象，需要创建
    final long preResolveHost = System.nanoTime();
    final InetSocketAddress resolvedAddress = new InetSocketAddress(remoteHost, remotePort);
    final long hostResolveTimeMs = (System.nanoTime() - preResolveHost) / 1000000;
    final String resolvMsg = resolvedAddress.isUnresolved() ? "failed" : "succeed";
    if (hostResolveTimeMs > 2000) {
      logger.warn("DNS resolution {} for {} took {} ms",
          resolvMsg, resolvedAddress, hostResolveTimeMs);
    } else {
      logger.trace("DNS resolution {} for {} took {} ms",
          resolvMsg, resolvedAddress, hostResolveTimeMs);
    }

    /**
     * ClientPool中的锁起作用的时候到了
     */
    synchronized (clientPool.locks[clientIndex]) {
      cachedClient = clientPool.clients[clientIndex];
      // 这里有点像 double-check，这样才正确，两个线程同时检测到没有TransportClient，第一个创建后第二个开始执行，他此时可能位于临界区了已经
      if (cachedClient != null) {
        if (cachedClient.isActive()) {
          logger.trace("Returning cached connection to {}: {}", resolvedAddress, cachedClient);
          return cachedClient;
        } else {
          logger.info("Found inactive connection to {}, creating a new one.", resolvedAddress);
        }
      }
      clientPool.clients[clientIndex] = createClient(resolvedAddress);
      return clientPool.clients[clientIndex];
    }
  }

  /**
   * Create a completely new {@link TransportClient} to the given remote host / port.
   * This connection is not pooled.
   *
   * As with {@link #createClient(String, int)}, this method is blocking.
   */
  public TransportClient createUnmanagedClient(String remoteHost, int remotePort)
      throws IOException, InterruptedException {
    final InetSocketAddress address = new InetSocketAddress(remoteHost, remotePort);
    return createClient(address);
  }

  /** Create a completely new {@link TransportClient} to the remote address.
   * 创建TransportClient的真正实现。
   */
  private TransportClient createClient(InetSocketAddress address)
      throws IOException, InterruptedException {
    logger.debug("Creating new connection to {}", address);

    Bootstrap bootstrap = new Bootstrap();
    bootstrap.group(workerGroup)
      .channel(socketChannelClass)
      // Disable Nagle's Algorithm since we don't want packets to wait
      .option(ChannelOption.TCP_NODELAY, true)
      .option(ChannelOption.SO_KEEPALIVE, true)
      .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectionTimeoutMs())
      .option(ChannelOption.ALLOCATOR, pooledAllocator);


    final AtomicReference<TransportClient> clientRef = new AtomicReference<>();
    final AtomicReference<Channel> channelRef = new AtomicReference<>();
    // 为根引导程序设置管道初始化回调函数，参考Netty中的知识
    // 此处的SocketChannel就是我们平时调用future.channel().writeAndFlush()时的channel
    // 所以真正写数据的时候依赖于它，必须保存好
    bootstrap.handler(new ChannelInitializer<SocketChannel>() {
      @Override
      public void initChannel(SocketChannel ch) {
        // 而TransportClient的真正创建过程是在该函数中
        // 该函数中创建了TransportClient并保存了最重要的channel
        TransportClient client = initializePipeline(ch, new TransportClientHandler());
        clientRef.set(client);
        channelRef.set(ch);
      }
    });

    // Connect to the remote server
    long preConnect = System.nanoTime();
    ChannelFuture cf = bootstrap.connect(address);
    if (!cf.await(connectionTimeoutMs())) {
      throw new IOException(
        String.format("Connecting to %s timed out (%s ms)", address, connectionTimeoutMs()));
    } else if (cf.cause() != null) {
      throw new IOException(String.format("Failed to connect to %s", address), cf.cause());
    }

    TransportClient client = clientRef.get();
    Channel channel = channelRef.get();
    assert client != null : "Channel future completed successfully with null client";
    return client;
  }


  /** Close all connections in the connection pool, and shutdown the worker thread pool. */
  @Override
  public void close() {
    // Go through all clients and close them if they are active.
    for (ClientPool clientPool : connectionPool.values()) {
      for (int i = 0; i < clientPool.clients.length; i++) {
        TransportClient client = clientPool.clients[i];
        if (client != null) {
          clientPool.clients[i] = null;
          JavaUtils.closeQuietly(client);
        }
      }
    }
    connectionPool.clear();

    if (workerGroup != null && !workerGroup.isShuttingDown()) {
      workerGroup.shutdownGracefully();
    }
  }

  private int connectionTimeoutMs() {
    return 10_000;
  }

  private TransportClient initializePipeline(SocketChannel channel, TransportClientHandler channelHandler) {
    try {
      /**
       * 对pipeline进行设置，这里的执行顺序参考Netty
       */
      ChannelPipeline pipeline = channel.pipeline()
              .addLast("encoder", ENCODER)
              .addLast("decoder", DECODER)
              .addLast("idleStateHandler",
                      new IdleStateHandler(0, 0, 10))
              .addLast("handler", channelHandler);
      TransportClient client = new TransportClient(channel);
      return client;
    } catch (RuntimeException e) {
      logger.error("Error while initializing Netty pipeline", e);
      throw e;
    }
  }
}
